# 1 MCP客户端
import asyncio

import chromadb
from langchain_community.embeddings import ModelScopeEmbeddings
from langchain_core.tools import BaseTool
from langchain_mcp_adapters.client import MultiServerMCPClient



async def main():
    mcp_client = MultiServerMCPClient(
        {
            "wms-scesrv": {
                "url": "http://47.102.147.217:28012/scesrv/sse",
                "transport": "sse",
                "headers": {
                    "X-Auth-token": "eef7TvPm5XsZ31CJW7uKAQOvP6kL6RQp"
                }
            },
            "wms-wmssrv": {
                "url": "http://47.102.147.217:30052/wmssrv/sse",
                "transport": "sse",
                "headers": {
                    "X-Auth-token": "eef7TvPm5XsZ31CJW7uKAQOvP6kL6RQp"
                }
            },
        }
    )

    # mcp_client = MultiServerMCPClient(
    #     {
    #         "flux-mcp": {
    #             "url": "http://127.0.0.1:9000/mcp",
    #             "transport": "streamable_http"
    #         }
    #     }
    # )

    # 1 得到用户操作
    # 2 向量数据库检索(根据用户操作需求查找最相近的几条) 前提:已经进行向量化了
    # 3 拼接prompt 筛选工具给Agent
    # 4 agent执行

    tools: list[BaseTool] = await mcp_client.get_tools(server_name="wms-scesrv")
    print(f"Get tools: {tools}")

    ids = []
    documents = []
    meta_datas = []

    for tool in tools:
        ids.append(tool.name)  # 唯一 id
        documents.append(tool.description or "")  # 将 description
        # metadata 里放你关心的字段, 比如 args_schema 等
        meta_datas.append({
            "args_schema": str(getattr(tool, "args_schema", None)),
        })

    model_id = "./nlp_gte_sentence-embedding_chinese-base"
    embeddings = ModelScopeEmbeddings(model_id=model_id)

    client = chromadb.Client()
    customer_support_collection = client.create_collection(
        name="tools"
    )

    customer_support_collection.add(
        ids=ids,
        documents=documents,
        metadatas=meta_datas,
    )

    user_query = "帮我新增一个系统代码"

    context = customer_support_collection.query(
        query_texts=[user_query],
        n_results=2
    )

    print(context["ids"])
    print(len(context["ids"][0]))


if __name__ == '__main__':
    asyncio.run(main())
